﻿/*
* Copyright (c) 2010-2012 Tesla Engine Group
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/

using System;
using System.IO;
using Assimp.Configs;
using Tesla.Core;
using Tesla.Graphics;
using Tesla.Math;
using Tesla.Scene;
using Tesla.Util;

namespace Tesla.Content.Loaders {

    /// <summary>
    /// Translator utility for taking a loaded Assimp scene and importing that into Tesla's default scene graph.
    /// </summary>
    public sealed class AssimpTranslator {

        private Assimp.AssimpImporter m_importer;
        private ContentManager m_content;
        private Matrix m_transform;

        private MeshData[] m_meshData;
        private Material[] m_materials;
        private MeshMaterialInfo[] m_meshMatInfo;

        private RasterizerState m_wireframeOneSidedRS;
        private RasterizerState m_twoSidedRS;
        private ModelLoaderParameters m_loaderParams;
        private IResource m_modelResource;

        public AssimpTranslator(ContentManager content) {
            m_content = content;
            m_importer = new Assimp.AssimpImporter();
        }

        public Spatial Load(IResource resource, String format, ModelLoaderParameters loaderParams) {

            if (!m_importer.IsFormatSupported(format)) {
                throw new TeslaException("Format not supported " + resource.Name);
            }

            m_modelResource = resource;
            m_loaderParams = loaderParams;
            m_importer.Scale = loaderParams.Scale;
            m_importer.XAxisRotation = loaderParams.XAxisRotation;
            m_importer.YAxisRotation = loaderParams.YAxisRotation;
            m_importer.ZAxisRotation = loaderParams.ZAxisRotation;

            Assimp.Scene scene = LoadFromAssimp(resource, format);

            if (scene == null) {
                throw new TeslaException("Could not load resource: " + resource.Name);
            }

            //TODO: Maybe return a not-found mesh rather than exception? May need to add that functionality to the content manager.

            ParseMeshData(scene);

            Spatial graph = ParseNodes(scene.RootNode);

            if (graph == null) {
                throw new TeslaException("Parse exception, no scene graph to load!" + resource.Name);
            }

            if (m_loaderParams.ImportLights) {
                ParseLights(scene, graph);
            }

            return graph;
        }

        private Assimp.Scene LoadFromAssimp(IResource resource, String format) {
            Assimp.PostProcessSteps ppFlags = Assimp.PostProcessPreset.TargetRealTimeMaximumQuality;

            switch (m_loaderParams.NormalGeneration) {
                //Preset default is to calc smooth normals, remove if none is set
                case NormalGeneration.None:
                    ppFlags ^= Assimp.PostProcessSteps.GenerateSmoothNormals;
                    break;
                //Remove smooth, set face
                case NormalGeneration.Face:
                    ppFlags ^= Assimp.PostProcessSteps.GenerateSmoothNormals;
                    ppFlags |= Assimp.PostProcessSteps.GenerateNormals;
                    //Per to assimp docs, since we need to dupe verts, not sure if this step runs, but just in case so not to waste time
                    ppFlags ^= Assimp.PostProcessSteps.JoinIdenticalVertices;
                    break;
                //If crease, we want to set the config setting otherwise we use the default assimp.
                case NormalGeneration.Crease:
                    m_importer.SetConfig(new NormalSmoothingAngleConfig(m_loaderParams.CreaseAngle));
                    break;

            }

            //Default in the preset is to calculate tangent basis, if we don't want to, remove it
            if (!m_loaderParams.GenerateTangentBasis) {
                ppFlags ^= Assimp.PostProcessSteps.CalculateTangentSpace;
            }

            //Flip UVs
            if (m_loaderParams.FlipUVCoordinates) {
                ppFlags |= Assimp.PostProcessSteps.FlipUVs;
            }

            //Default winding order in assimp is CCW. In tesla it's CW, so false we flip, true we dont
            if (!m_loaderParams.SwapWindingOrder) {
                ppFlags |= Assimp.PostProcessSteps.FlipWindingOrder;
            }

            using (Stream str = resource.OpenStream()) {
                return m_importer.ImportFileFromStream(str, ppFlags, format);
            }
        }

        private Spatial ParseNodes(Assimp.Node assimpNode, Node parent = null) {
            if (!assimpNode.HasMeshes && !assimpNode.HasChildren) {
                return null;
            }

            if (!assimpNode.HasChildren && assimpNode.HasMeshes) {
                Spatial spatial;
                if (assimpNode.MeshCount == 1) {
                    spatial = CreateMesh(assimpNode);
                } else {
                    spatial = CreateNode(assimpNode);
                }
                if (parent != null) {
                    parent.AddChild(spatial);
                }
                return spatial;
            }

            Node node = CreateNode(assimpNode);

            foreach (Assimp.Node child in assimpNode.Children) {
                ParseNodes(child, node);
            }

            if (parent != null) {
                parent.AddChild(node);
            }

            return node;
        }

        private Node CreateNode(Assimp.Node assimpNode) {
            Node node = new Node(assimpNode.Name);
            node.Transform.Set(ConvertMatrix(assimpNode.Transform));

            if (assimpNode.HasMeshes) {
                int[] indices = assimpNode.MeshIndices;
                for(int i = 0; i < indices.Length; i++) {
                    int meshIndex = indices[i];
                    Mesh child = new Mesh(String.Format("{0}.mesh-part{1}", node.Name, i));
                    child.MeshData = GetMeshData(meshIndex);
                    child.Material = GetMaterial(GetMaterialIndex(meshIndex));
                    node.AddChild(child);
                }
            }
            return node;
        }

        private Mesh CreateMesh(Assimp.Node assimpNode) {
            if (!assimpNode.HasMeshes) {
                return null;
            }

            Mesh mesh = new Mesh(assimpNode.Name);
            mesh.Transform.Set(ConvertMatrix(assimpNode.Transform));

            int meshIndex = assimpNode.MeshIndices[0];
            mesh.MeshData = GetMeshData(meshIndex);
            mesh.Material = GetMaterial(GetMaterialIndex(meshIndex));

            return mesh;
        }

        private void ParseLights(Assimp.Scene scene, Spatial root) {
            if (!scene.HasLights) {
                return;
            }

            Assimp.Light[] lights = scene.Lights;
            Node isNode = root as Node;

            for (int i = 0; i < lights.Length; i++) {
                Assimp.Light light = lights[i];
                String name = light.Name;
                Spatial spatial = root;

                if (!String.IsNullOrEmpty(name)) {
                    if (isNode != null) {
                        Spatial descendent = isNode.GetDescendentNamed(name);
                        if (descendent != null) {
                            spatial = descendent;
                        }
                    }
                }

                switch (light.LightType) {
                    case Assimp.LightSourceType.Spot:
                        CreateSpotLight(light, spatial);
                        break;
                    case Assimp.LightSourceType.Point:
                        CreatePointLight(light, spatial);
                        break;
                    case Assimp.LightSourceType.Directional:
                        CreateDirectionalLight(light, spatial);
                        break;
                }
            }
        }

        private void CreatePointLight(Assimp.Light light, Spatial spatial) {
            PointLight pl = new PointLight();
            pl.Constant = light.AttenuationConstant;
            pl.Linear = light.AttenuationLinear;
            pl.Quadratic = light.AttenuationQuadratic;
            pl.Ambient = ConvertColor(light.ColorAmbient);
            pl.Diffuse = ConvertColor(light.ColorDiffuse);
            pl.Specular = ConvertColor(light.ColorSpecular);
            pl.Position = ConvertVector3(light.Position);
            spatial.AddLight(pl);
        }

        private void CreateDirectionalLight(Assimp.Light light, Spatial spatial) {
            DirectionalLight dl = new DirectionalLight();
            dl.Ambient = ConvertColor(light.ColorAmbient);
            dl.Diffuse = ConvertColor(light.ColorDiffuse);
            dl.Specular = ConvertColor(light.ColorSpecular);
            dl.Direction = ConvertVector3(light.Direction);
            spatial.AddLight(dl);
        }

        private void CreateSpotLight(Assimp.Light light, Spatial spatial) {
            SpotLight sl = new SpotLight();
            sl.Constant = light.AttenuationConstant;
            sl.Linear = light.AttenuationLinear;
            sl.Quadratic = light.AttenuationQuadratic;
            sl.Ambient = ConvertColor(light.ColorAmbient);
            sl.Diffuse = ConvertColor(light.ColorDiffuse);
            sl.Specular = ConvertColor(light.ColorSpecular);
            sl.Position = ConvertVector3(light.Position);
            sl.Direction = ConvertVector3(light.Direction);
            spatial.AddLight(sl);
        }

        private MeshData GetMeshData(int index) {
            if (index >= 0 || index < m_meshData.Length) {
                return m_meshData[index];
            }
            return null;
        }

        private Material GetMaterial(int index) {
            if (index >= 0 || index < m_materials.Length) {
                return m_materials[index];
            }
            return null;
        }

        private int GetMaterialIndex(int meshIndex) {
            if (meshIndex >= 0 || meshIndex < m_meshMatInfo.Length) {
                return m_meshMatInfo[meshIndex].MaterialIndex;
            }
            return -1;
        }

        private bool ParseMeshData(Assimp.Scene scene) {
            if (!scene.HasMeshes) {
                return false;
            }

            Assimp.Mesh[] meshes = scene.Meshes;
            if (meshes.Length == 0) {
                return false;
            }

            m_meshData = new MeshData[meshes.Length];
            m_meshMatInfo = new MeshMaterialInfo[meshes.Length];

            for (int i = 0; i < meshes.Length; i++) {
                Assimp.Mesh mesh = meshes[i];
                MeshData md = new MeshData();
                MeshMaterialInfo info = new MeshMaterialInfo();

                if (mesh.HasVertices) {
                    Assimp.Vector3D[] assimpPos = mesh.Vertices;
                    DataBuffer<Vector3> positions = new DataBuffer<Vector3>(assimpPos.Length);
                    MemoryHelper.Copy(assimpPos, positions.Buffer, Vector3.SizeInBytes * assimpPos.Length);
                    md.Positions = positions;
                }

                if (mesh.HasNormals) {
                    Assimp.Vector3D[] assimpNorms = mesh.Normals;
                    DataBuffer<Vector3> normals = new DataBuffer<Vector3>(assimpNorms.Length);
                    MemoryHelper.Copy(assimpNorms, normals.Buffer, Vector3.SizeInBytes * assimpNorms.Length);
                    md.Normals = normals;
                    info.HasNormals = true;
                }

                if (mesh.HasTangentBasis) {
                    Assimp.Vector3D[] assimpTangs = mesh.Tangents;
                    Assimp.Vector3D[] assimpBitangs = mesh.BiTangents;
                    DataBuffer<Vector3> tangents = new DataBuffer<Vector3>(assimpTangs.Length);
                    DataBuffer<Vector3> bitangents = new DataBuffer<Vector3>(assimpBitangs.Length);
                    MemoryHelper.Copy(assimpTangs, tangents.Buffer, Vector3.SizeInBytes * assimpTangs.Length);
                    MemoryHelper.Copy(assimpBitangs, bitangents.Buffer, Vector3.SizeInBytes * assimpBitangs.Length);
                    md.Tangents = tangents;
                    md.Binormals = bitangents;
                    info.HasTangentBasis = true;
                }

                if (mesh.HasVertexColors(0)) {
                    Assimp.Color4D[] assimpColors = mesh.GetVertexColors(0);
                    md.Colors = new DataBuffer<Color>(ConvertColors(assimpColors));
                    info.HasVertexColors = true;
                }

                if (mesh.HasTextureCoords(0)) {
                    Assimp.Vector3D[] assimpUV = mesh.GetTextureCoords(0);
                    md.TextureCoordinates = new DataBuffer<Vector2>(ConvertVectors(assimpUV));
                    info.HasTexCoords = true;
                }

                md.Indices = new DataBuffer<int>(mesh.GetIntIndices());
                md.Reconstruct();
                m_meshMatInfo[i] = info;
                m_meshData[i] = md;
            }

            ParseMaterials(scene);
            return true;
        }

        private void ParseMaterials(Assimp.Scene scene) {
            Assimp.Material[] materials = scene.Materials;

            //Should always have the default material at least
            m_materials = new Material[materials.Length];
            bool hasNormals = false;
            bool hasTangentBasis = false;
            bool hasTexCoords = false;
            bool hasVertexColors = false;

            //TODO: Incorporate user-defined materials

            for (int i = 0; i < materials.Length; i++) {
                Assimp.Material material = materials[i];
                Material mat;
                hasNormals = false;
                hasTangentBasis = false;
                hasTexCoords = false;
                hasVertexColors = false;

                foreach (MeshMaterialInfo info in m_meshMatInfo) {
                    if (info.MaterialIndex == i) {
                        if (info.HasNormals) {
                            hasNormals = true;
                        }
                        if (info.HasTangentBasis) {
                            hasTangentBasis = true;
                        }
                        if (info.HasTexCoords) {
                            hasTexCoords = true;
                        }
                        if (info.HasVertexColors) {
                            hasVertexColors = true;
                        }
                    }
                }

                if (hasNormals && m_loaderParams.PreferLitMaterials) {
                    if (hasTexCoords && hasTangentBasis) {
                        if (hasVertexColors) {
                            mat = m_content.Load<Material>("NormalMapVertColor.tem").Clone();
                        } else {
                            mat = m_content.Load<Material>("NormalMap.tem").Clone();
                        }
                        SetNormalMapMaterialTextures(mat, material); //Normal map textures + diffuse
                    } else {
                        if (hasTexCoords) {
                            if (hasVertexColors) {
                                mat = m_content.Load<Material>("LitBasicVertColorTexture.tem");
                            } else {
                                mat = m_content.Load<Material>("LitBasicTexture.tem").Clone();
                            }
                            SetMaterialTexture(mat, material); //Difuse texture
                        } else {
                            if (hasVertexColors) {
                                mat = m_content.Load<Material>("LitBasicVertColor.tem").Clone();
                            } else {
                                mat = m_content.Load<Material>("LitBasicColor.tem").Clone();
                            }
                        }
                    }
                    SetLitMaterialColors(mat, material); //Common lit color parameters
                } else {
                    if (hasTexCoords) {
                        if (hasVertexColors) {
                            mat = m_content.Load<Material>("BasicVertColorTexture.tem");
                        } else {
                            mat = m_content.Load<Material>("BasicTexture.tem").Clone();
                        }
                        SetMaterialTexture(mat, material); //Difuse texture
                    } else {
                        if (hasVertexColors) {
                            mat = m_content.Load<Material>("BasicVertColor.tem").Clone();
                        } else {
                            mat = m_content.Load<Material>("BasicColor.tem").Clone();
                        }
                    }
                    SetMaterialColor(mat, material); //Common unlit color parameters
                }

                SetMaterialRenderStates(mat, material);

                if (material.HasName) {
                    mat.Name = material.Name;
                }
                m_materials[i] = mat;
            }
        }

        //Per-model import cleanup
        private void Cleanup() {
            m_materials = null;
            m_meshData = null;
            m_meshMatInfo = null;
            m_loaderParams = null;
            m_importer.DetachLogStreams();
            m_importer.RemoveConfig(NormalSmoothingAngleConfig.NormalSmoothingAngleConfigName);
        }

        private void SetLitMaterialColors(Material material, Assimp.Material assimpMat) {
            if (assimpMat.HasColorAmbient) {
                material.SetParameter("MatAmbient", ConvertColorToVector3(assimpMat.ColorAmbient));
            }

            if (assimpMat.HasColorDiffuse) {
                material.SetParameter("MatDiffuse", ConvertColorToVector3(assimpMat.ColorDiffuse));
            }

            if (assimpMat.HasColorEmissive) {
                material.SetParameter("MatEmissive", ConvertColorToVector3(assimpMat.ColorEmissive));
            }

            if (assimpMat.HasColorSpecular) {
                material.SetParameter("MatSpecular", ConvertColorToVector3(assimpMat.ColorSpecular));
            }

            if (assimpMat.HasShininessStrength) {
                material.SetParameter("Shininess", assimpMat.ShininessStrength);
            }

            if (assimpMat.HasOpacity) {
                material.SetParameter("Alpha", assimpMat.Opacity);
            }
        }

        private void SetMaterialColor(Material material, Assimp.Material assimpMat) {
            if (assimpMat.HasColorDiffuse) {
                material.SetParameter("DiffuseColor", ConvertColor(assimpMat.ColorDiffuse));
            }
        }

        private void SetMaterialTexture(Material material, Assimp.Material assimpMat) {
            Assimp.TextureSlot diffuseMap = assimpMat.GetTexture(Assimp.TextureType.Diffuse, 0);
            if (!String.IsNullOrEmpty(diffuseMap.FilePath)) {
                Texture diffuse;
                try {
                    diffuse = m_content.LoadRelativeTo<Texture2D>(diffuseMap.FilePath, m_modelResource);// m_content.Load<Texture2D>(diffuseMap.FilePath);
                    material.SetParameter("DiffuseMap", diffuse);
                } catch {
                //    diffuse = m_content.Load<Texture2D>("notexture.tebo");
                }
            }
        }

        private void SetNormalMapMaterialTextures(Material material, Assimp.Material assimpMat) {
            Assimp.TextureSlot diffuseMap = assimpMat.GetTexture(Assimp.TextureType.Diffuse, 0);
            Assimp.TextureSlot normalMap = assimpMat.GetTexture(Assimp.TextureType.Normals, 0);
            Assimp.TextureSlot specularMap = assimpMat.GetTexture(Assimp.TextureType.Specular, 0);

            if (!String.IsNullOrEmpty(diffuseMap.FilePath)) {
                Texture diffuse;
                try {
                    diffuse = m_content.LoadRelativeTo<Texture2D>(diffuseMap.FilePath, m_modelResource); //m_content.Load<Texture2D>(diffuseMap.FilePath);
                    material.SetParameter("DiffuseMap", diffuse);
                } catch {
                 //   diffuse = m_content.Load<Texture2D>("notexture.tebo");
                }
            }

            if (!String.IsNullOrEmpty(normalMap.FilePath)) {
                try {
                    Texture normal = m_content.LoadRelativeTo<Texture2D>(normalMap.FilePath, m_modelResource); //m_content.Load<Texture2D>(normalMap.FilePath);
                    material.SetParameter("NormalMap", normal);
                } catch {
                    //Default normal map?
                }
            }

            if (!String.IsNullOrEmpty(specularMap.FilePath)) {
                try {
                    Texture specular = m_content.LoadRelativeTo<Texture2D>(specularMap.FilePath, m_modelResource);//m_content.Load<Texture2D>(specularMap.FilePath);
                    material.SetParameter("SpecularMap", specular);
                } catch {
                    //Default specular map?
                }
            }
        }

        private void SetMaterialRenderStates(Material material, Assimp.Material assimpMat) {

            if (assimpMat.HasTwoSided && assimpMat.HasWireFrame) {
               // material.SetRenderState(RasterizerState.CullNoneWireframe);
            } else if (assimpMat.HasWireFrame) {
                if (m_wireframeOneSidedRS == null) {
                    m_wireframeOneSidedRS = new RasterizerState();
                    m_wireframeOneSidedRS.Fill = FillMode.WireFrame;
                    m_wireframeOneSidedRS.BindRenderState();
                }
               // material.SetRenderState(m_wireframeOneSidedRS);
            } else if (assimpMat.HasTwoSided) {
                if (m_twoSidedRS == null) {
                    m_twoSidedRS = new RasterizerState();
                    m_twoSidedRS.Cull = CullMode.None;
                    m_twoSidedRS.BindRenderState();
                }
              //  material.SetRenderState(m_twoSidedRS);
            } 

            if (assimpMat.HasBlendMode) {
                if (assimpMat.BlendMode == Assimp.BlendMode.Additive) {
                  //  material.SetRenderState(BlendState.AdditiveBlend);
                } else {
                  //  material.SetRenderState(BlendState.AlphaBlendNonPremultiplied);
                }
            }
        }

        private Color[] ConvertColors(Assimp.Color4D[] colors) {
            Color[] data = new Color[colors.Length];
            for (int i = 0; i < colors.Length; i++) {
                Assimp.Color4D color = colors[i];
                data[i] = new Color(color.R, color.G, color.B, color.A);
            }
            return data;
        }

        private Vector2[] ConvertVectors(Assimp.Vector3D[] vectors) {
            Vector2[] data = new Vector2[vectors.Length];
            for (int i = 0; i < vectors.Length; i++) {
                Assimp.Vector3D vector = vectors[i];
                data[i] = new Vector2(vector.X, 1 -  vector.Y);
            }
            return data;
        }

        private Vector3 ConvertVector3(Assimp.Vector3D vector) {
            return new Vector3(vector.X, vector.Y, vector.Z);
        }

        private Quaternion ConvertQuaternion(Assimp.Quaternion quat) {
            return new Quaternion(quat.X, quat.Y, quat.Z, quat.W);
        }

        private Matrix ConvertMatrix(Assimp.Matrix4x4 matrix) {
            return new Matrix(matrix.A1, matrix.B1, matrix.C1, matrix.D1, matrix.A2, matrix.B2, matrix.C2, matrix.D2, matrix.A3, matrix.B3, matrix.C3, matrix.D3,
                matrix.A4, matrix.B4, matrix.C4, matrix.D4);
        }

        private Color ConvertColor(Assimp.Color4D color) {
            return new Color(color.R, color.G, color.B, color.A);
        }

        private Color ConvertColor(Assimp.Color3D color) {
            return new Color(color.R, color.G, color.B);
        }

        private Vector3 ConvertColorToVector3(Assimp.Color4D color) {
            return new Vector3(color.R, color.G, color.B);
        }

        private void CreateModelTransform() {
            Matrix scale = Matrix.FromScale(m_loaderParams.Scale);
            Matrix xRot = Matrix.FromRotationX(m_loaderParams.XAxisRotation);
            Matrix yRot = Matrix.FromRotationY(m_loaderParams.YAxisRotation);
            Matrix zRot = Matrix.FromRotationZ(m_loaderParams.ZAxisRotation);

            m_transform = (scale) * (xRot * yRot * zRot);
        }

        private struct MeshMaterialInfo {
            public int MaterialIndex;
            public bool HasNormals;
            public bool HasTangentBasis;
            public bool HasTexCoords;
            public bool HasVertexColors;

            public MeshMaterialInfo(int matIndex, bool hasNormals, bool hasTangentBasis, bool hasTexCoords, bool hasVertexColors) {
                MaterialIndex = matIndex;
                HasNormals = hasNormals;
                HasTangentBasis = hasTangentBasis;
                HasTexCoords = hasTexCoords;
                HasVertexColors = hasVertexColors;
            }
        }
    }
}
